import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt

# 创建示例数据集
data = {
    'text': [
        'win free prize money now',
        'urgent meeting tomorrow office',
        'special discount only today',
        'hello how are you doing',
        'claim your lottery ticket',
        'lunch at 12pm tomorrow',
        'buy now get free shipping',
        'see you later my friend',
        'limited time offer buy',
        'project deadline next week'
    ],
    'label': ['spam', 'ham', 'spam', 'ham', 'spam', 'ham', 'spam', 'ham', 'spam', 'ham']
}

df = pd.DataFrame(data)
print("数据集:")
print(df)

# 特征提取：将文本转换为词频向量
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(df['text'])
y = df['label']

print(f"\n特征矩阵形状: {X.shape}")
print(f"词汇表: {vectorizer.get_feature_names_out()}")

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 创建并训练朴素贝叶斯模型
model = MultinomialNB()
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

print(f"\n预测结果:")
print(f"真实标签: {list(y_test)}")
print(f"预测标签: {list(y_pred)}")

# 评估模型
accuracy = accuracy_score(y_test, y_pred)
print(f"\n准确率: {accuracy:.2f}")
print(f"\n分类报告:")
print(classification_report(y_test, y_pred))

# 可视化混淆矩阵
plt.figure(figsize=(8, 6))
cm = confusion_matrix(y_test, y_pred)
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
            xticklabels=['ham', 'spam'], 
            yticklabels=['ham', 'spam'])
plt.title('Confusion Matrix')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.show()

# 测试新数据
new_emails = [
    'free money offer today',
    'meeting schedule tomorrow',
    'win lottery prize now'
]

new_X = vectorizer.transform(new_emails)
predictions = model.predict(new_X)
probabilities = model.predict_proba(new_X)

print("\n新邮件预测:")
for email, pred, prob in zip(new_emails, predictions, probabilities):
    print(f"邮件: '{email}'")
    print(f"预测: {pred}, 概率: [ham: {prob[0]:.3f}, spam: {prob[1]:.3f}]")
    print()